# -*- coding: utf-8 -*-
"""
emoji规则，在模型预测的基础上根据表情修改标签
"""
import os
import json
import numpy as np
import pandas as pd
import emoji

# 表情规则
rule_dict = {
    "😊": "happy",
    "😍": "happy",
    "💕": "happy",
    "😄": "happy",
    "😜": "happy",
    "😰": "sad",
    "😪": "sad",
    "😢": "sad",
    "😤": "angry"
}


def load_eval_data(is_virus: bool, data_folder=r'../raw/eval'):
    """
    加载通用和疫情验证集
    :param is_virus: usual or virus
    :param data_folder: 数据文件夹
    :return:
    """
    # usual_eval.txt
    if is_virus:
        # virus_eval.txt
        path = os.path.join(data_folder, 'virus_eval.txt')
    else:
        path = os.path.join(data_folder, 'usual_eval.txt')

    print('Loading {}...'.format(path))
    with open(path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    df = pd.DataFrame()
    df['id'] = [i['id'] for i in data]
    df['content'] = [i['content'] for i in data]

    return df


def generate_result(is_virus: bool, predict_prob: np.ndarray, label_dict_path, output_folder):
    """
    :param is_virus: usual(False) / virus(True)
    :param predict_prob: 模型输出概率(softmax)，shape: (N, 6)
    :param label_dict_path: 标签转换
    :param output_folder: 输出文件夹
    :return:
    """
    df = load_eval_data(is_virus=is_virus)
    assert len(df) == predict_prob.shape[0]
    # print(predict_prob)

    # load label dict
    with open(label_dict_path, 'r', encoding='utf-8') as f:
        label_dict = json.load(f)
    label_dict = {value: key for key, value in label_dict.items()}

    result = []
    change_count = 0  # 统计有多少数据被修改
    for index, content, prob in zip(df['id'].astype(int), df['content'].astype(str), predict_prob):
        if emoji.emoji_count(content) > 0:
            # 表情规则
            count = emoji.emoji_count(content)
            for i in emoji.emoji_lis(content):
                if i['emoji'] in rule_dict:
                    result.append({"id": index, "label": rule_dict[i['emoji']]})
                    # 输出看看原始标签和被修改的标签
                    print('{} -> {},{}'.format(label_dict[np.argmax(prob)], rule_dict[i['emoji']], content))
                    change_count += 1
                    break
                else:
                    count -= 1
            if count == 0:
                # 句子中所有表情均不在规则中
                result.append({"id": index, "label": label_dict[np.argmax(prob)]})
        else:
            # 句子没有表情，取概率最大的标签输出
            result.append({"id": index, "label": label_dict[np.argmax(prob)]})

    assert len(result) == len(df)

    # output
    path = os.path.join(output_folder, '{}_result.txt'.format('virus' if is_virus else 'usual'))
    with open(path, 'w', encoding='utf-8') as f:
        json.dump(result, f)
    print('Done.Change count:', change_count)


def generate_result2(is_virus: bool, predict, label_dict_path, output_folder):
    """
    :param is_virus: usual(False) / virus(True)
    :param predict: 模型输出结果，[{"id": 0, "label": ""}]
    :param label_dict_path: 标签转换
    :param output_folder: 输出文件夹
    :return:
    """
    df = load_eval_data(is_virus=is_virus)
    assert len(df) == len(predict)
    predict = [i['label'] for i in predict]
    # print(predict_prob)

    # load label dict
    with open(label_dict_path, 'r', encoding='utf-8') as f:
        label_dict = json.load(f)
    label_dict = {value: key for key, value in label_dict.items()}

    result = []
    change_count = 0  # 统计有多少数据被修改
    for index, content, predict_label in zip(df['id'].astype(int), df['content'].astype(str), predict):
        if emoji.emoji_count(content) > 0:
            # 表情规则
            count = emoji.emoji_count(content)
            for i in emoji.emoji_lis(content):
                if i['emoji'] in rule_dict:
                    result.append({"id": index, "label": rule_dict[i['emoji']]})
                    # 输出看看原始标签和被修改的标签
                    print('id: {}, {} -> {},{}'.format(index, predict_label, rule_dict[i['emoji']], content))
                    change_count += 1
                    break
                else:
                    count -= 1
            if count == 0:
                # 句子中所有表情均不在规则中
                result.append({"id": index, "label": predict_label})
        else:
            # 句子没有表情，取原始标签输出
            result.append({"id": index, "label": predict_label})

    assert len(result) == len(df)

    # output
    path = os.path.join(output_folder, '{}_result.txt'.format('virus' if is_virus else 'usual'))
    with open(path, 'w', encoding='utf-8') as f:
        json.dump(result, f)
    print('Done.Change count:', change_count)


if __name__ == '__main__':
    import torch
    # generate_result(is_virus=True,
    #                 predict_prob=torch.softmax(torch.rand(2000, 6), dim=1).numpy(),
    #                 label_dict_path=r'../data/label_dict.json',
    #                 output_folder=r'./')

    # generate_result(is_virus=False,
    #                 predict_prob=torch.softmax(torch.rand(2000, 6), dim=1).numpy(),
    #                 label_dict_path=r'../data/label_dict.json',
    #                 output_folder=r'./')

    # -------------------------------------------------------------------------------

    # 小心会覆盖掉原来的
    # with open(r'./usual_result.txt', 'r', encoding='utf-8') as f:
    #     predict_file = json.load(f)
    # generate_result2(is_virus=False,
    #                  predict=predict_file,
    #                  label_dict_path=r'../data/label_dict.json',
    #                  output_folder=r'./')
